package com.raos.example.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.noear.snack.core.utils.StringUtil;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * ollama工具类
 *
 * @author raos
 * @date 2025-02-13
 */
@Slf4j
public class OllamaUtil {

    public static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";
    public static final String DEFAULT_DEEPSEEK_MODEL_NAME = "deepseek-r1:1.5b";

    /**
     * 创建回复文本
     *
     * @param modelName 模型名
     * @param prompt 提示词
     * @return
     */
    public static String generateText(String modelName, String prompt) {
        try {
            if (StringUtil.isEmpty(modelName) || StringUtil.isEmpty(modelName.trim())) {
                modelName = DEFAULT_DEEPSEEK_MODEL_NAME;
            }
            URL url = new URL(OLLAMA_API_URL);
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("POST");
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setDoOutput(true);

            String jsonInputString = "{\"model\": \""+ modelName +"\", \"prompt\": \"" + prompt + "\"}";

            try (OutputStream os = connection.getOutputStream()) {
                byte[] input = jsonInputString.getBytes(StandardCharsets.UTF_8);
                os.write(input, 0, input.length);
            }

            ObjectMapper objectMapper = new ObjectMapper(); // JSON 序列化工具
            StringBuilder response = new StringBuilder();
            try (InputStreamReader reader = new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8);
                 BufferedReader br = new BufferedReader(reader)) {
                String responseLine;
                while ((responseLine = br.readLine()) != null) {
                    String str = responseLine.trim();
                    log.debug("responseLine：[{}]", str);
                    // response.append(str);
                    Map map = objectMapper.readValue(str, Map.class);
                    if (map.containsKey("response") && map.get("response") != null) {
                        response.append(map.get("response").toString());
                    }
                }
            }

            return response.toString();
        } catch (Exception e) {
            log.error("获取Ollama回复文本异常", e);
            return "Error: " + e.getMessage();
        }
    }

}
